﻿using System.Text.Json;
using ModelContextProtocol.Server;
using Microsoft.SemanticKernel.Memory;
using Microsoft.SemanticKernel.Connectors.Chroma;
using Microsoft.Extensions.Logging;

namespace edan.zhjk.Commons.MCPService.ChromaDB
{
    /// <summary>
    /// Chroma
    /// </summary>
    [McpServerToolType]

    public class ChromaDBTool
    {
        //private readonly ILogger<ChromaDBTool> _logger;
        //private readonly ChromaMemoryStore _memoryStore;

        //public ChromaDBTool(ILogger<ChromaDBTool> logger, ChromaMemoryStore memoryStore)
        //{
        //    _logger = logger;
        //    _memoryStore = memoryStore;
        //}

        //// 模型类
        //public class TextContent
        //{
        //    public string Type { get; set; } = "text";
        //    public string Text { get; set; } = "";
        //}

        //public class ToolUsageItem
        //{
        //    public string Name { get; set; } = "";
        //    public Dictionary<string, object> Args { get; set; } = new Dictionary<string, object>();
        //}

        //public class LogChatInput
        //{
        //    public string PromptSummary { get; set; } = "";
        //    public string ResponseSummary { get; set; } = "";
        //    public string RawPrompt { get; set; } = "";
        //    public string RawResponse { get; set; } = "";
        //    public List<ToolUsageItem> ToolUsage { get; set; } = new List<ToolUsageItem>();
        //    public List<Dictionary<string, object>> FileChanges { get; set; } = new List<Dictionary<string, object>>();
        //    public string InvolvedEntities { get; set; } = "";
        //    public string SessionId { get; set; } = "";
        //    public string CollectionName { get; set; } = "chat_history_v1";
        //}

        //public class DocumentMetadata
        //{
        //    public string CollectionName { get; set; } = "";
        //    public string Document { get; set; } = "";
        //    public string Id { get; set; } = "";
        //    public Dictionary<string, object> Metadata { get; set; } = new Dictionary<string, object>();
        //}

        ///// <summary>
        ///// 日志记录实现
        ///// </summary>
        ///// <param name="input"></param>
        ///// <returns></returns>
        //public async Task<List<TextContent>> LogChatAsync(LogChatInput input)
        //{
        //    try
        //    {
        //        var chatId = Guid.NewGuid().ToString();
        //        var metadata = new Dictionary<string, object>
        //        {
        //            ["prompt_summary"] = input.PromptSummary,
        //            ["response_summary"] = input.ResponseSummary,
        //            ["raw_prompt"] = input.RawPrompt,
        //            ["raw_response"] = input.RawResponse,
        //            ["tool_usage"] = JsonSerializer.Serialize(input.ToolUsage),
        //            ["file_changes"] = JsonSerializer.Serialize(input.FileChanges),
        //            ["involved_entities"] = input.InvolvedEntities,
        //            ["session_id"] = input.SessionId,
        //            ["timestamp"] = DateTimeOffset.UtcNow.ToUnixTimeSeconds()
        //        };

        //        // 确保集合存在
        //        if (!await _memoryStore.DoesCollectionExistAsync(input.CollectionName))
        //        {
        //            await _memoryStore.CreateCollectionAsync(input.CollectionName);
        //        }

        //        // 创建内存记录
        //        var memoryRecord = MemoryRecord.LocalRecord(
        //            id: chatId,
        //            text: $"{input.PromptSummary}\n{input.ResponseSummary}",
        //            description: "Chat log",
        //            additionalMetadata: JsonSerializer.Serialize(metadata));

        //        // 保存到 ChromaDB
        //        await _memoryStore.UpsertAsync(input.CollectionName, memoryRecord);

        //        return new List<TextContent>
        //        {
        //            new TextContent { Text = JsonSerializer.Serialize(new { success = true, chat_id = chatId }) }
        //        };
        //    }
        //    catch (Exception ex)
        //    {
        //        _logger.LogError(ex, "Error logging chat");
        //        return new List<TextContent>
        //        {
        //            new TextContent { Text = JsonSerializer.Serialize(new { success = false, error = ex.Message }) }
        //        };
        //    }
        //}

        ///// <summary>
        ///// 文档操作实现
        ///// </summary>
        ///// <param name="input"></param>
        ///// <returns></returns>
        //public async Task<List<TextContent>> AddDocumentAsync(DocumentMetadata input)
        //{
        //    try
        //    {
        //        // 确保集合存在
        //        if (!await _memoryStore.DoesCollectionExistAsync(input.CollectionName))
        //        {
        //            await _memoryStore.CreateCollectionAsync(input.CollectionName);
        //        }

        //        // 处理元数据中的时间戳
        //        input.Metadata = EnsureServerTimestamp(input.Metadata);

        //        // 创建内存记录
        //        var memoryRecord = MemoryRecord.LocalRecord(
        //            id: string.IsNullOrEmpty(input.Id) ? Guid.NewGuid().ToString() : input.Id,
        //            text: input.Document,
        //            additionalMetadata: JsonSerializer.Serialize(input.Metadata));

        //        // 保存到 ChromaDB
        //        var id = await _memoryStore.UpsertAsync(input.CollectionName, memoryRecord);

        //        return new List<TextContent>
        //        {
        //            new TextContent { Text = JsonSerializer.Serialize(new { added_id = id }) }
        //        };
        //    }
        //    catch (Exception ex)
        //    {
        //        _logger.LogError(ex, "Error adding document");
        //        return new List<TextContent>
        //        {
        //            new TextContent { Text = JsonSerializer.Serialize(new { error = ex.Message }) }
        //        };
        //    }
        //}

        //public async Task<List<TextContent>> QueryDocumentsAsync(string collectionName, List<string> query, bool withEmbeddings = false)
        //{
        //    try
        //    {
        //        var results = new List<Dictionary<string, object>>();

        //        // 执行查询
        //        var memories = _memoryStore.GetBatchAsync(
        //            collectionName,
        //            query,
        //            withEmbeddings: withEmbeddings);

        //        await foreach (var memory in memories)
        //        {
        //            var metadata = JsonSerializer.Deserialize<Dictionary<string, object>>(
        //                memory.Metadata.AdditionalMetadata ?? "{}");

        //            results.Add(new Dictionary<string, object>
        //            {
        //                ["id"] = memory.Metadata.Id,
        //                ["text"] = memory.Metadata.Text,
        //                ["metadata"] = metadata,
        //                ["embedding"] = memory.Embedding
        //            });
        //        }

        //        return new List<TextContent>
        //        {
        //            new TextContent { Text = JsonSerializer.Serialize(results) }
        //        };
        //    }
        //    catch (Exception ex)
        //    {
        //        _logger.LogError(ex, "Error querying documents");
        //        return new List<TextContent>
        //        {
        //            new TextContent { Text = JsonSerializer.Serialize(new { error = ex.Message }) }
        //        };
        //    }
        //}

        ///// <summary>
        ///// 集合操作实现
        ///// </summary>
        ///// <param name="collectionName"></param>
        ///// <returns></returns>
        //public async Task<List<TextContent>> CreateCollectionAsync(string collectionName)
        //{
        //    try
        //    {
        //        await _memoryStore.CreateCollectionAsync(collectionName);
        //        return new List<TextContent> {
        //        new TextContent { Text = $"Collection '{collectionName}' created successfully" }
        //    };
        //    }
        //    catch (Exception ex)
        //    {
        //        _logger.LogError(ex, "Error creating collection");
        //        return new List<TextContent>
        //        {
        //            new TextContent { Text = JsonSerializer.Serialize(new { error = ex.Message }) }
        //        };
        //    }
        //}

        //public async Task<List<TextContent>> ListCollectionsAsync(int limit = 10, int offset = 0)
        //{
        //    try
        //    {
        //        var collections = await _memoryStore.GetCollectionsAsync().ToListAsync();
        //        var result = new
        //        {
        //            collections = collections.Skip(offset).Take(limit),
        //            total = collections.Count
        //        };

        //        return new List<TextContent> {
        //        new TextContent { Text = JsonSerializer.Serialize(result) }
        //    };
        //    }
        //    catch (Exception ex)
        //    {
        //        _logger.LogError(ex, "Error listing collections");
        //        return new List<TextContent>
        //        {
        //            new TextContent { Text = JsonSerializer.Serialize(new { error = ex.Message }) }
        //        };
        //    }
        //}

        ///// <summary>
        ///// 辅助方法
        ///// </summary>
        ///// <param name="metadata"></param>
        ///// <returns></returns>
        //private Dictionary<string, object> EnsureServerTimestamp(Dictionary<string, object> metadata)
        //{
        //    var timestampFields = new[] { "timestamp", "created_at", "updated_at" };
        //    var now = DateTimeOffset.UtcNow;

        //    foreach (var field in timestampFields)
        //    {
        //        if (metadata.ContainsKey(field))
        //        {
        //            metadata[field] = now.ToUnixTimeSeconds();
        //        }
        //    }

        //    return metadata;
        //}

        //        private const string CollectionId = "fake-collection-id";
        //        private const string CollectionName = "fake-collection-name";

        //        private readonly HttpClient _httpClient;
        //#pragma warning disable SKEXP0020 // 类型仅用于评估，在将来的更新中可能会被更改或删除。取消此诊断以继续。
        //        private readonly IChromaClient _chromaClient;
        //        public ChromaDBTool()
        //        {

        //        }


        //        public async Task<MemoryRecord?> GetAsync()
        //        {
        //            // Arrange
        //            const string Endpoint = "https://fake-random-test-host/fake-path/";

        //            var store = new ChromaMemoryStore(this._httpClient, Endpoint);

        //            return await store.GetAsync("fake-collection", "fake-key");
        //        }


        //        public async Task<MemoryRecord?> GetAsync()
        //        {
        //            var store = GetChromaMemoryStoreByHttpClient();

        //            return await store.GetAsync("fake-collection", "fake-key");
        //        }

        //        /// <summary>
        //        /// 创建Chroma
        //        /// </summary>
        //        /// <returns></returns>
        //        public async Task CreateCollectionAsync()
        //        {
        //            var store = GetChromaMemoryStore();
        //            await store.CreateCollectionAsync(CollectionName);
        //        }


        //        public async Task DeleteCollectionAsync()
        //        {
        //            var store = GetChromaMemoryStore();
        //            await store.DeleteCollectionAsync(CollectionName);
        //        }

        //        /// <summary>
        //        /// 确定数据存储中是否存在某一集合
        //        /// </summary>
        //        /// <returns></returns>
        //        public async Task<bool> DoesCollectionExistAsync()
        //        {
        //            var store = GetChromaMemoryStore();
        //            return await store.DoesCollectionExistAsync(CollectionName);
        //        }


        //        public async Task<MemoryRecord?> GetAsync()
        //        {
        //            // Arrange
        //            var expectedMemoryRecord = this.GetRandomMemoryRecord();
        //            var embeddingsModel = this.GetEmbeddingsModelFromMemoryRecord(expectedMemoryRecord);

        //            this._chromaClient.GetEmbeddingsAsync(CollectionId, new[] { expectedMemoryRecord.Key }, It.IsAny<string[]>(), CancellationToken.None)
        //                .ReturnsAsync(embeddingsModel);

        //            var store = GetChromaMemoryStore();
        //            return await store.GetAsync(CollectionName, expectedMemoryRecord.Key, withEmbedding: true);
        //        }


        //        public async Task<MemoryRecord?> ItReturnsNullWhenMemoryRecordDoesNotExistAsync()
        //        {
        //            // Arrange
        //            const string MemoryRecordKey = "fake-record-key";

        //            _chromaClient.GetEmbeddingsAsync(CollectionId, new[] { MemoryRecordKey }, It.IsAny<string[]>(), CancellationToken.None);

        //            var store = GetChromaMemoryStore();
        //            return await store.GetAsync(CollectionName, MemoryRecordKey, withEmbedding: true);
        //        }


        //        public async Task<MemoryRecord> ItThrowsExceptionOnGettingMemoryRecordFromNonExistingCollectionAsync()
        //        {
        //            // Arrange
        //            const string CollectionName = "non-existent-collection";
        //            const string MemoryRecordKey = "fake-record-key";
        //            const string CollectionDoesNotExistErrorMessage = $"Collection {CollectionName} does not exist";

        //            await this._chromaClient.GetCollectionAsync(CollectionName, CancellationToken.None);

        //            var store = GetChromaMemoryStore();
        //            return await store.GetAsync(CollectionName, MemoryRecordKey, withEmbedding: true);
        //        }


        //        public async IAsyncEnumerable<MemoryRecord> GetBatchAsync()
        //        {
        //            MemoryRecord[] expectedMemoryRecords = [memoryRecord1, memoryRecord2, memoryRecord3];
        //            var memoryRecordKeys = expectedMemoryRecords.Select(l => l.Key).ToArray();

        //            var embeddingsModel = this.GetEmbeddingsModelFromMemoryRecords(expectedMemoryRecords);

        //            _chromaClient.GetEmbeddingsAsync(CollectionId, memoryRecordKeys, It.IsAny<string[]>(), CancellationToken.None);

        //            var store = GetChromaMemoryStore();
        //            store.GetBatchAsync(CollectionName, memoryRecordKeys, withEmbeddings: true);
        //        }


        //        public IAsyncEnumerable<string> GetCollectionsAsync()
        //        {
        //            _chromaClient.ListCollectionsAsync(CancellationToken.None);

        //            var store = GetChromaMemoryStore();
        //            store.GetCollectionsAsync();
        //        }

        //        #region private ================================================================================

        //        /// <summary>
        //        /// 为存储在内部存储提供程序中的内存准备一个实例
        //        /// </summary>
        //        /// <param name="embedding">源内容嵌入</param>
        //        /// <returns></returns>
        //        private MemoryRecord GetLocalRecord(ReadOnlyMemory<float>? embedding = null)
        //        {
        //            var id = Guid.NewGuid().ToString();
        //            var memoryEmbedding = embedding ?? new[] { 1f, 3f, 5f };

        //            return MemoryRecord.LocalRecord(
        //                id: id, //存储提供者中的资源标识符，例如记录ID/GUID/增量计数器等。
        //                text: "text-" + Guid.NewGuid().ToString(),  //全文用于生成嵌入
        //                description: "description-" + Guid.NewGuid().ToString(),  //可选的记录描述。注意：描述没有索引。
        //                embedding: memoryEmbedding, //源内容嵌入
        //                additionalMetadata: "metadata-" + Guid.NewGuid().ToString(),    //可选字符串，用于保存自定义元数据。
        //                key: id);   //可选的现有数据库密钥
        //        }

        //        /// <summary>
        //        /// 
        //        /// </summary>
        //        /// <param name="memoryRecord"></param>
        //        /// <returns></returns>
        //        private Dictionary<string, object> GetMemoryRecordDictionary(MemoryRecord memoryRecord)
        //        {
        //            var serialized = JsonSerializer.Serialize(memoryRecord.Metadata);
        //            return JsonSerializer.Deserialize<Dictionary<string, object>>(serialized)!;
        //        }

        //        /// <summary>
        //        /// 
        //        /// </summary>
        //        /// <param name="memoryRecords"></param>
        //        /// <returns></returns>
        //        private ChromaEmbeddingsModel GetEmbeddingsModelFromMemoryRecords(MemoryRecord[] memoryRecords)
        //        {
        //            var embeddingsModel = new ChromaEmbeddingsModel();

        //            embeddingsModel.Ids.AddRange(memoryRecords.Select(l => l.Key));
        //            embeddingsModel.Embeddings.AddRange(memoryRecords.Select(l => l.Embedding.ToArray()));
        //            embeddingsModel.Metadatas.AddRange(memoryRecords.Select(this.GetMemoryRecordDictionary));

        //            return embeddingsModel;
        //        }

        //        /// <summary>
        //        /// 
        //        /// </summary>
        //        /// <param name="memoryRecord"></param>
        //        /// <returns></returns>
        //        private ChromaEmbeddingsModel GetEmbeddingsModelFromMemoryRecords(MemoryRecord memoryRecord)
        //        {
        //            return this.GetEmbeddingsModelFromMemoryRecords([memoryRecord]);
        //        }

        //        #endregion warning restore SKEXP0020 // 类型仅用于评估，在将来的更新中可能会被更改或删除。取消此诊断以继续。

        //        #pragma warning disable SKEXP0020 // 类型仅用于评估，在将来的更新中可能会被更改或删除。取消此诊断以继续。
        //        /// <summary>
        //        /// 
        //        /// </summary>
        //        /// <returns></returns>
        //        private ChromaMemoryStore GetChromaMemoryStoreByHttpClient()
        //        {
        //            // Arrange
        //            const string BaseAddress = "https://fake-random-test-host/fake-path/";

        //            using var httpClient = this.GetHttpClientStub();
        //            httpClient.BaseAddress = new Uri(BaseAddress);

        //            return new ChromaMemoryStore(httpClient);
        //        }

        //        /// <summary>
        //        /// 
        //        /// </summary>
        //        /// <returns></returns>
        //        public ChromaMemoryStore GetChromaMemoryStore()
        //        {
        //            return new ChromaMemoryStore(_chromaClient);
        //        }
        //        #endregion warning disable SKEXP0020 // 类型仅用于评估，在将来的更新中可能会被更改或删除。取消此诊断以继续。

    }

}
